﻿using System.IO;
using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using UnityEditor;
using System.Reflection;
using UnityEngine.Rendering;
using UnityEngine.Profiling;

namespace Jelly.Editor.AssetBundle
{
    /// <summary>
    /// 收集shader的shader_feature keyword到svc的工具类，防止在bundle中丢失变种，可在构建时自动收集，或手动扫描构建资源收集
    /// 用法：1. 创建一个shaderVariantCollection文件，并使其进入构建
    ///      2. 对进入构建的每个材质文件调用OnMaterialScanned
    ///      3. 分析依赖之后调用UpdateShaderVariantCollection
    ///      4. 构建
    /// 原理： 1. 通过进入构建的材质找到keywords和shader
    ///       2. 所有扫描完成后，用反射执行内部函数，找到shader的所有passType和keywords列表
    ///       3. 把匹配的存到svc中，svc进入构建
    ///       4. （更完善的方法）基于扫描出来的svc，创建自定义数据，运行时可分帧warmup
    /// 注意shader_feature_local_fragment的keywrod没法通过svc的UI加入，改成shader_feature_local就可以了
    ///    multi_compile的keyword不必收集，必然进入bundle，除非被IPreprocessShaders剔除
    /// </summary>
    public class ShaderVariantCollector
    {
        class ShaderInfoOnMaterial
        {
            public Shader shader;
            // 每个string是多个keyword，为了去重，没用数组
            public HashSet<string> keywordSeqSet = new HashSet<string>();
        }


        //一个shader的变种们，一个变种包括passtype和keyword列表
        public class ShaderData
        {
            public Shader shader;
            public int[] passtypes;
            // 每个string是多个keyword，它的下标和passtypes是对应的
            public string[] keywordSeqList;
        }


        Dictionary<Shader, ShaderInfoOnMaterial> s_ShaderInfoOnMaterial;
        HashSet<string> s_ScannedMatPath;

        System.Diagnostics.Stopwatch s_StopWatch;
        long s_ScanMaterialTicks = 0;
        long s_GenVaraintMS = 0;
        long s_SaveSVCAssetMS = 0;
        List<ShaderVariantCollection.ShaderVariant> s_ShaderVaraintsUsedByMaterial;
        ShaderVariantCollection s_TempSVC;

        //基于反射的方法来查找shader的keywords
        MethodInfo s_GetShaderVariantEntries = null;
        MethodInfo s_AddNewShaderToCollection = null;


        /// <summary>
        /// 开始执行前需要重置
        /// </summary>
        public ShaderVariantCollector()
        {
            s_ShaderInfoOnMaterial = new Dictionary<Shader, ShaderInfoOnMaterial>();
            s_ScannedMatPath = new HashSet<string>();
            s_ScanMaterialTicks = 0;
            s_GenVaraintMS = 0;
            s_SaveSVCAssetMS = 0;
            s_ShaderVaraintsUsedByMaterial = new List<ShaderVariantCollection.ShaderVariant>();
            s_TempSVC = new ShaderVariantCollection();
            s_StopWatch = System.Diagnostics.Stopwatch.StartNew();
            s_StopWatch.Stop();

            s_GetShaderVariantEntries = typeof(ShaderUtil).GetMethod("GetShaderVariantEntriesFiltered",
                BindingFlags.NonPublic | BindingFlags.Static);
            s_AddNewShaderToCollection = typeof(ShaderUtil).GetMethod("AddNewShaderToCollection",
                BindingFlags.NonPublic | BindingFlags.Static);
        }

        /// <summary>
        /// 扫描新材质时的回调
        /// </summary>
        /// <param name="matPath">材质路径</param>
        public void OnMaterialScanned(string matPath)
        {
            Profiler.BeginSample("OnMaterialScanned");
            s_StopWatch.Restart();
            if (string.IsNullOrEmpty(matPath) || s_ScannedMatPath.Contains(matPath))
            {
                s_ScanMaterialTicks += s_StopWatch.ElapsedTicks;
                Profiler.EndSample();
                return;
            }

            Material mat = AssetDatabase.LoadAssetAtPath<Material>(matPath);
            if (mat == null)
            {
                Debug.LogError($"OnMaterialScanned: mat is null, matPath:{matPath}");
            }
            else if(mat.shader == null)
            {
                Debug.LogError($"OnMaterialScanned: mat.shader is null, matPath:{matPath}");
            }
            
            if (!s_ShaderInfoOnMaterial.TryGetValue(mat.shader, out ShaderInfoOnMaterial info))
            {
                info = new ShaderInfoOnMaterial() { shader = mat.shader };
                s_ShaderInfoOnMaterial.Add(info.shader, info);
            }
            string keywordSeq = string.Join(" ", mat.shaderKeywords);
            bool added = info.keywordSeqSet.Add(keywordSeq);
            s_ScannedMatPath.Add(matPath);
            s_ScanMaterialTicks += s_StopWatch.ElapsedTicks;
            Profiler.EndSample();

            // Debug.LogWarning($"keywords added:{added}, {keywordSeq}, matPath:{matPath}");
        }

        /// <summary>
        /// 扫描后更新svc
        /// </summary>
        public void UpdateShaderVariantCollection(bool saveAssets)
        {
            Profiler.BeginSample("GenVariant");
            s_StopWatch.Restart();
            foreach (var kv in s_ShaderInfoOnMaterial)
            {
                FindShaderVariantsUsedByMaterial(kv.Value);
            }
            s_GenVaraintMS = s_StopWatch.ElapsedMilliseconds;
            Profiler.EndSample();


            // 变种计入到svc中
            Profiler.BeginSample("save svc");
            s_StopWatch.Restart();
            var svcPath = AssetBuilderCfg.Instance.svcPath;
            var svcAsset = AssetDatabase.LoadAssetAtPath<ShaderVariantCollection>(svcPath);
            if (svcAsset == null)
            {
                svcAsset = new ShaderVariantCollection();
                Directory.CreateDirectory(Path.GetDirectoryName(svcPath));
                AssetDatabase.CreateAsset(svcAsset, svcPath);
            }

            int shaderInc = svcAsset.shaderCount;
            int varaintInc = svcAsset.variantCount;
            var oldSVCShaders = FindShaderInSVC(svcAsset);
            var variantsAdded = new List<ShaderVariantCollection.ShaderVariant>();
            foreach (var v in s_ShaderVaraintsUsedByMaterial)
            {
                if (!svcAsset.Contains(v))
                    variantsAdded.Add(v);

                svcAsset.Add(v);
            }
            shaderInc = svcAsset.shaderCount - shaderInc;
            varaintInc = svcAsset.variantCount - varaintInc;

            EditorUtility.SetDirty(svcAsset);
            // 内存中已经有数据了，不SaveAssets也能在构建中识别
            if (saveAssets)
                AssetDatabase.SaveAssets();
            s_SaveSVCAssetMS = s_StopWatch.ElapsedMilliseconds;
            Profiler.EndSample();

            int kewordSeqCount = 0;
            foreach (var v in s_ShaderInfoOnMaterial)
                kewordSeqCount += v.Value.keywordSeqSet.Count;

            var scanMatTime = s_ScanMaterialTicks * 1f / System.Diagnostics.Stopwatch.Frequency;
            var sb = new System.Text.StringBuilder();
            sb.AppendLine($"ShaderVariantCollecter report:");
            sb.AppendLine($"【scan】{scanMatTime:F3}s, material: {s_ScannedMatPath.Count}, shader:{s_ShaderInfoOnMaterial.Count}, unique shaderKeywords: {kewordSeqCount}");
            sb.AppendLine($"【gen 】{s_GenVaraintMS * 1e-3}s, shaderVariant: {s_ShaderVaraintsUsedByMaterial.Count}");
            sb.AppendLine($"【svc 】{s_SaveSVCAssetMS * 1e-3}s, add new shader: {shaderInc}, new variant: {varaintInc}");
            variantsAdded.Sort((x, y) => { return x.shader.name.CompareTo(y.shader.name); });
            bool csvStyle = false;
            if (csvStyle)
            {
                sb.AppendLine("id, newShader, shaderName, filePath, keywords, passType");
                for (int i = 0; i < variantsAdded.Count; ++i)
                {
                    var sv = variantsAdded[i];
                    var status = !oldSVCShaders.Contains(sv.shader) ? "1" : "0";
                    var filePath = AssetDatabase.GetAssetPath(sv.shader);
                    sb.AppendLine($"  {i}, {status}, {sv.shader.name}, {filePath}, {string.Join(" ", sv.keywords)}, {sv.passType}");
                }
            }
            else
            {
                for (int i = 0; i < variantsAdded.Count; ++i)
                {
                    var sv = variantsAdded[i];
                    var status = !oldSVCShaders.Contains(sv.shader) ? "newShader+++" : "newVariant";
                    sb.AppendLine($"  {i}, {status}, shaderName: {sv.shader.name}, keywords: {string.Join(" ", sv.keywords)}, passType: {sv.passType}");
                }
            }
            Debug.LogWarning(sb.ToString());
        }

        public HashSet<Shader> FindShaderInSVC(ShaderVariantCollection svc)
        {
            SerializedObject so = new SerializedObject(svc);
            SerializedProperty m_Shaders = so.FindProperty("m_Shaders");
            HashSet<Shader> result = new HashSet<Shader>();

            for (var i = 0; i < m_Shaders.arraySize; ++i)
            {
                var entryProp = m_Shaders.GetArrayElementAtIndex(i);
                Shader shader = (Shader)entryProp.FindPropertyRelative("first").objectReferenceValue;
                if (shader != null)
                    result.Add(shader);
            }
            return result;
        }



        List<ShaderVariantCollection.ShaderVariant> FindShaderVariantsUsedByMaterial(ShaderInfoOnMaterial sim)
        {
            foreach (var keywordSeq in sim.keywordSeqSet)
            {
                // Debug.LogWarning(keywordSeq);
                int candidateCount = 2;     // 候选是包含keywordSeq的，在后悬中找跟keywordSeq一致的，可能有多个pass？
                var (sd, firstMatch) = GetShaderDataByKeywordSeq(sim.shader, candidateCount, keywordSeq);
                // Debug.LogWarning($"shader all variants count:{sd.keywordSeqList.Length}");
                if (firstMatch)
                {
                    for (int i = 0; i < sd.keywordSeqList.Length; ++i)
                    {
                        // Debug.Log($"idnex:{i}, passType:{(PassType)sd.passtypes[i]}, keywords:{sd.keywordSeqList[i]}");
                        if (keywordSeq == sd.keywordSeqList[i])
                        {
                            var sv = new ShaderVariantCollection.ShaderVariant(sd.shader, (PassType)sd.passtypes[i], keywordSeq.Split(' '));
                            s_ShaderVaraintsUsedByMaterial.Add(sv);
                            string logKeyword = string.IsNullOrEmpty(keywordSeq) ? "<no keyword>" : keywordSeq;
                            // Debug.LogWarning($"add {sd.shader.name}, keywords: {logKeyword}");
                        }
                    }
                }
                else
                {
                    // 没有一摸一样的，处理掉无效keyword试一试
                    (sd, firstMatch) = GetShaderDataByKeywordSeq(sim.shader, 1, keywordSeq, true);
                    if ((firstMatch || !AssetBuilderCfg.Instance.shaderVariantCollectorAddExactMatch) && 
                        sd.keywordSeqList != null && sd.keywordSeqList.Length > 0)
                    {
                        string newkeywordSeq = sd.keywordSeqList[0];
                        var sv = new ShaderVariantCollection.ShaderVariant(sd.shader, (PassType)sd.passtypes[0], newkeywordSeq.Split(' '));
                        s_ShaderVaraintsUsedByMaterial.Add(sv);
                        string logKeyword = string.IsNullOrEmpty(newkeywordSeq) ? "<no keyword>" : newkeywordSeq;
                        if (!firstMatch)
                            Debug.LogWarning($"shader {sim.shader.name} not contians keywords:{keywordSeq}, add similar one:{logKeyword}");
                    }
                    else
                        Debug.LogWarning($"shader {sim.shader.name} not contians keywords:{keywordSeq}, no similar one is added");
                }
            }
            return s_ShaderVaraintsUsedByMaterial;
        }


        /// <summary>
        /// 通过反射找到shader的keywrods信息
        /// @Note 对于理论变种数较少的情况，可以得到所有可能(selectedKeywords传空)，再按每个材质的keywords搜索
        ///       但对于理论变种数极大的情况，可行的做法是传入材质中的keyworlds，在返回值中找一样kewords的。每种材质keywords都调用此函数
        /// </summary>
        /// <param name="shader"></param>
        /// <param name="count">返回信息的最大条数</param>
        /// <param name="selectedKeywords">返回信息中必须要包含的keyword列表，用空格分隔，可以是""字符串</param>
        /// <param name="testKeywordValidity">true时测试传入keyword是否有效，无效的会被去掉，再进行查询</param>
        /// <returns>每个keywords string看似都是字符串排序过的，空格分隔</returns>
        (ShaderData, bool) GetShaderDataByKeywordSeq(Shader shader, int count, string selectedKeywords = "",
                           bool testKeywordValidity = false)
        {
            // 2019.4要先把shader加入到svc中，否则后面获取变种会崩溃
            s_TempSVC.Clear();
            s_AddNewShaderToCollection.Invoke(null, new object[] { shader, s_TempSVC });
            // 基于反射获取shader所有的type和keyword
            // internal static void GetShaderVariantEntriesFiltered(Shader shader, int maxEntries, string[] filterKeywords, 
            //          ShaderVariantCollection excludeCollection, out int[] passTypes, out string[] keywordLists, out string[] remainingKeywords)
            // 向svc资源通过其UI添加keywords，就是用GetShaderVariantEntriesFiltered来获取可选择的列表，通过已选的keyword来过滤
            // filterKeywords是svc界面中用户选择的某些keyword，用于缩小进一步选择范围
            // excludeCollection是svc，返回的keywordLists会排除已经加入到那个svc中的
            // passTypes是所有变种的pass类型
            // keywordLists是包括filterKeywords时的每个变种的keywords列表
            // remainingKeywords是所有的单个keyword，因为svc中并没有选择keyword，全是剩余的

            int[] passTypes = new int[count];
            string[] keywordLists = new string[] { };
            string[] remainingKeywords = new string[] { };
            // "".split(' ')没法找到无keyword的情况
            string[] filter = string.IsNullOrEmpty(selectedKeywords) ? new string[] { } : selectedKeywords.Split(' ');
            object[] args;

            // 去掉无效keyword
            if (testKeywordValidity && filter.Length > 0)
            {
                args = new object[] { shader, 1, new string[] { }, s_TempSVC, passTypes, keywordLists, remainingKeywords };
                s_GetShaderVariantEntries.Invoke(null, args);
                remainingKeywords = args[6] as string[];
                var newFilter = new List<string>();
                foreach (var v in filter)
                {
                    if (ArrayUtility.Contains<string>(remainingKeywords, v))
                        newFilter.Add(v);
                }
                filter = newFilter.ToArray();
            }

            args = new object[] { shader, count, filter, s_TempSVC, passTypes, keywordLists, remainingKeywords };
            s_GetShaderVariantEntries.Invoke(null, args);
            passTypes = args[4] as int[];
            keywordLists = args[5] as string[];

            ShaderData sdata = new ShaderData() { shader = shader, passtypes = passTypes, keywordSeqList = keywordLists };
            bool firstOneMatch = false;     // 返回的第一个是否和传入的（是否去除无效都行）完全一致
            if (keywordLists.Length > 0 && keywordLists[0] == string.Join(" ", filter))
            {
                firstOneMatch = true;
            }
            return (sdata, firstOneMatch);
        }
        
        /// <summary>
        /// 遍历扫描过的材质，检查其中是否有开启instancing的，如果开启，则创建临时材质启用instancing
        /// 避免INSTANCING_ON被裁剪
        /// </summary>
        /// <param name="shaderParentDir">shader目录路径</param>
        /// <param name="folderName">创建材质的目录名</param>
        /// <returns>被添加instancing材质的shader集合</returns>
        public HashSet<Shader> AddInstancingMaterial(string shaderParentDir, string folderName = "TempMat")
        {
            s_StopWatch.Restart();
            shaderParentDir = shaderParentDir.Replace("\\", "/");
            var tempMatSaveDirPath = shaderParentDir + "/" + folderName;
            if (Directory.Exists(tempMatSaveDirPath))
            {
                Debug.Log($"删除 {tempMatSaveDirPath} 目录，重新生成instancing材质");
                AssetDatabase.DeleteAsset(tempMatSaveDirPath);
            }
            AssetDatabase.CreateFolder(shaderParentDir, folderName);

            // 已经添加过临时instancing材质的shader的集合
            var addedMatShaderSet = new HashSet<Shader>();
            foreach (var materialPath in s_ScannedMatPath)
            {
                var mat = AssetDatabase.LoadAssetAtPath<Material>(materialPath);
                var shader = mat.shader;
                if (addedMatShaderSet.Contains(shader))
                {
                    continue;
                }
                var addAutoMat = false;
                if (mat.enableInstancing)
                {
                    var tempMat = new Material(shader);
                    tempMat.enableInstancing = true;
                    AssetDatabase.CreateAsset(tempMat, tempMatSaveDirPath + "/" + shader.name.Replace("/","_") + ".mat");
                    addAutoMat = true;
                }
                if (addAutoMat)
                {
                    addedMatShaderSet.Add(shader);
                }
            }
            Debug.Log($"添加instancing材质完成，耗时{s_StopWatch.ElapsedMilliseconds * 1e-3}s，共有 {addedMatShaderSet.Count} shader添加instancing材质");
            return addedMatShaderSet;
        }
    }

}